#pragma once

#include "concurrent_queue.hpp"
#include <iostream>
#include <atomic>
#include <memory>
#include <functional>
#include <future>
#include <vector>

// 基于 concurrent_queue , 模拟 ctpl 来实现一个线程池
namespace ctpl {
    class ctpl_thread_pool {
    private:
    // 默认 10 个线程
    #define DFL_THREAD_NUM 10
        using self = ctpl_thread_pool;
    public:
        ctpl_thread_pool(int n_threads = DFL_THREAD_NUM) {
            init();
            resize(n_threads);
        }

        ~ctpl_thread_pool() { stop(true); }

        int size()                              { return static_cast<int>(_threads.size()); }

        int n_idle()                            { return static_cast<int>(_n_idle); }

        std::thread & get_thread(int i)         { return *_threads[i]; }

        // 这里的调用场景应该是在单线程下
        void resize(int n_threads) {
            if (_is_stop == false && _is_done == false) {
                int old_n_threads = _threads.size();
                if (old_n_threads < n_threads) {            // 扩容
                    _threads.resize(n_threads);
                    _flags.resize(n_threads);

                    for (int i = old_n_threads; i < n_threads; i ++) {
                        _flags[i] = std::make_shared<std::atomic<bool>>(false);      // 初始化
                        set_thread(i);
                    }
                }
                else {                                      // 缩容
                    for (int i = n_threads; i < old_n_threads; i ++) {
                        *(_flags[i]) = true;                                // 给这些线程发送信号
                        _threads[i]->detach();                              // 并分离线程，过一会就会被回收
                    }
                    {
                        std::unique_lock<mutex> lock(_mtx);
                        _cv.notify_all();                   // 唤醒所有的线程，让被删除的线程退出工作
                    }

                    _flags.resize(n_threads);
                    _threads.resize(n_threads);
                }
            }
        }

        // clear all the tasks in queue
        void clear_queue() {
            std::function<void()> *f_p;
            while (_q.try_pop(f_p) == true) {       // false 表示为空
                delete f_p;
            }
        }

        // is_wait: 如果为 true，表示等待线程池中的任务执行完
        //          如果为 false，所有线程执行完当前任务之后就停止工作
        void stop(bool is_wait) {
            if (is_wait == false) {     // dont wait, force all threads to stop
                if (_is_stop)
                    return ;

                _is_stop = true;
                for (int i = 0; i < _threads.size(); i ++) {
                    *(_flags[i]) = true;                    // 标记为 true , 这些线程执行完任务后就会退出
                }

                clear_queue();
            }
            else {                      // wait all threads to finish all tasks
                if (_is_stop || _is_done)
                    return ;

                _is_done = true;        // 标记为 true , 完成线程池中所有的任务之后退出
            }

            {
                std::unique_lock<mutex> lock(_mtx); // 唤醒所有的线程
                _cv.notify_all();
            }

            // 开始清理工作
            for (int i = 0; i < _threads.size(); i ++) {
                if (_threads[i]->joinable()) {
                    _threads[i]->join();            // 回收资源
                }
            }
            // 如果这时候线程已经被回收了，但是队列里面仍然有任务，那么就会在这里被销毁
            clear_queue();
            _threads.clear();
            _flags.clear();
        }

        template<class F, class ...Rest>
        auto push(F && f, Rest &&... rest) -> std::future<decltype(f(rest...))> {
            // usage: packaged_task<int(int)>
            auto pck_p = make_shared<std::packaged_task<decltype(f(rest...))()>> (
                std::bind(std::forward<F>(f), std::forward<Rest>(rest)...)
                // []() {f(rest...);}
            );
            auto f_p = new function<void()>([pck_p](){
                (*pck_p)();
            });
            _q.push(f_p);
            std::unique_lock<mutex> lock(_mtx);
            _cv.notify_one();
            return pck_p->get_future();
        }

        // 往线程池里面添加一个新的任务, 会返回 future 对象，方便调用者获取线程的执行结果, 返回值就是任务函数的返回值, 也可以从里面获取异常
        template<class F>   
        auto push(F && f) -> std::future<decltype(f())> {
            // 用一个 function<void()> 来包装用户函数
            auto pck_p = make_shared<std::packaged_task<decltype(f())()>> (std::forward<F>(f));  // f 就是函数指针，直接给 shared_ptr 管理
            // 构建一个 function<void()> 包装用户函数，并放入安全队列中
            auto f_p = new function<void()>([pck_p]() {
                (*pck_p)();                             // 调用用户函数
            });
            _q.push(f_p);                               // push a task in queue
            std::unique_lock<mutex> lock(_mtx);
            _cv.notify_one();                               
            return pck_p->get_future();
        }

        // 从任务队列中取出一个原生的函数
        std::function<void()> pop() {
            std::function<void()> *f_p = nullptr;
            _q.pop(f_p);
            std::unique_ptr<std::function<void()>> uni_f(f_p);       // 自行管理
            std::function<void()> ret_f;
            if (f_p != nullptr) 
                ret_f = *f_p;
            return ret_f;
        }

    private:
        // 创建一个线程然后为这个线程作初始化
        void set_thread(int i) {
            // 重新为这个标志位设置一份拷贝
            // auto flag = std::make_shared<atomic<bool>> (_flags[i]);  // 这种写法错误
            std::shared_ptr<std::atomic<bool>> flag (_flags[i]);
            auto f = [flag /**a copy*/, i, this]() {
                std::cout << i << " thread begin to work" << std::endl;
                std::atomic<bool>& flag_ref = *flag;                         // get the ref of the flag
                std::function<void()> *f_p = nullptr;
                bool is_popped = _q.try_pop(f_p);
                while (true) {
                    // 会不断从队列中取出任务
                    while (is_popped == true) {         // if there is any task in queue
                        // 弹出来的函数指针用 unqiue_ptr 管理
                        std::unique_ptr<std::function<void()>> uni_f(f_p);
                        std::cout << "fetch one task, begin to run\n";
                        (*f_p)();        // run the task
                        if (flag_ref == true) {     // 执行任务的时候发现标志位被修改
                            std::cout << "current thread's flag is true, exit" << std::endl;
                            return ;                // over
                        }
                        else {
                            is_popped = _q.try_pop(f_p);
                        }
                    }
                    // 这里应该只是用于挂起线程，防止队列没任务，但是大量线程空转
                    std::unique_lock<mutex> lock(_mtx);
                    _n_idle ++;
                    _cv.wait(lock, [this, &is_popped, &f_p, &flag_ref]() { is_popped = _q.try_pop(f_p); return is_popped || _is_done || flag_ref; });
                    _n_idle --;
                    // 如果没弹出来，说明是 _is_done || flag_ref
                    if (is_popped == false) {
                        std::cout << "_is_done or current thread's flag is true, exit" << std::endl;
                        return ;
                    }
                }
            };
            // reset the managed resources for this unique_ptr (of cause need to be ptr)
            _threads[i].reset(new std::thread(f));
        }

    private:
        void init() {
            _is_stop        = false;
            _is_done        = false;
            _n_idle         = 0;
        }

        ctpl_thread_pool(const self&)   = delete;
        self& operator=(const self)     = delete;
        ctpl_thread_pool(self&&)        = delete;
        self& operator=(self&&)         = delete;

    private:
        std::vector<std::unique_ptr<std::thread>> _threads;        // 下标就是线程的标志位, 用来管理所有的线程
        std::vector<shared_ptr<std::atomic<bool>>> _flags;     // 用来控制所有的线程运行状态
        ConcurrentQueue<std::function<void()> *> _q;
        // _is_stop 是线程处理完一个任务之后，不管队列里面还有没有，马上停止工作
        std::atomic<bool> _is_stop;          // 该线程池是否暂停
        // 如果 _is_done 为 true，那么线程就尽量将所有的任务都处理完，再回收
        std::atomic<bool> _is_done;          // 全局性的标志位
        std::atomic<int> _n_idle;            // 正在空闲等待的线程数量

        mutex _mtx;                     // 用来保证多个线程对队列的同步操作
        condition_variable _cv;         
    };
}

